#include <net/sock.h>
#include <linux/module.h>
#include <linux/inet.h>
#include <linux/skbuff.h>
#include <linux/if_arp.h>
#include <linux/netdevice.h>
#include <linux/if.h>
#include <linux/if_vlan.h>
#include <linux/delay.h>
#include <net/udp_tunnel.h>
// #include <net/sch_generic.h>
#include <linux/kfifo.h>
#include <linux/netfilter.h>
#include <linux/semaphore.h>
static unsigned short DPORT = 17780, SPORT = 12345, SRC_PORT = 17780;
const static int sendedDataCount = 0, detectPktNum = 1000, min_delay = 0, max_delay = 2,min_delay1 = 0, max_delay1 = 2;

struct Recv_sockets
{
    struct socket *sk4;
    struct socket *sk6;
    struct socket *dataSkt;
} recv_sockets;

struct splb_info
{
    __be32 src_ip;
    __be32 dst_ip;
    __be16 src_port;
    __be16 dst_port;
    __be16 hon_id;
    __be16 option_port;
};

struct hon_info
{
    __u16 hon_id;
    __u16 src_port;
};
///////////////////////////////////////数据包队列
// struct dataPkt_queue_node
// {
//     //struct sk_buff *skb;
//     unsigned short port; //回传探测包携带的端口信息
//     struct dataPkt_queue_node *next;
// };

// static struct dataPkt_queue_node *head = NULL, *rear = NULL;

// struct semaphore skb_num;

// void initqueue(void)
// {
//     sema_init(&skb_num, 0);
//     head = kmalloc(sizeof(struct dataPkt_queue_node), 0);
//     if (head->next != NULL)
//         printk("可是head不是NULL了");
//     rear = head;
//     head->next = NULL;
//     rear->next = NULL;
// }

// int getqueueLength(void)
// {
//     int i = 0;
//     struct dataPkt_queue_node *tmp = head;
//     if (tmp->next == NULL)
//         return 0;
//     // if(head->next!=NULL) printk("有病吧！");
//     // printk("快点显示");
//     while (tmp->next != NULL)
//     {
//         tmp = tmp->next;
//         i++;
//     }
//     return i;
// }
// void enqueue(unsigned short port)
// {
//     struct dataPkt_queue_node *a;
//     a = kmalloc(sizeof(struct dataPkt_queue_node), 0);
//     a->port = port;
//     a->next = NULL;
//     rear->next = a;
//     rear = rear->next;
// }

// unsigned short dequeue(void)
// {
//     if (head->next == NULL)
//         return NULL;
//     // struct sk_buff *skb = head->next->skb;
//     unsigned short port = head->next->port;
//     struct dataPkt_queue_node *tmp = head->next;
//     head->next = tmp->next;
//     if (head->next == NULL)
//         head = rear;
//     if (tmp != NULL)
//     {
//         kfree(tmp);
//         printk("释放一个node!");
//     }
//     return port;
// }

// void drainqueue(void)
// {
//     int i = getqueueLength();
//     while (i > 0)
//     {
//         i--;
//         dequeue();
//     }
// }
#define BUFF_SIZE 4096                   //队列buffer大小
static char queue_buff[BUFF_SIZE] = {0}; //队列，用于存储回传探测包携带的端口号
static struct kfifo *pkfifo;

struct semaphore skb_num; //用于同步的信号量

void initqueue(void)
{
    sema_init(&skb_num, 0);
    pkfifo = kmalloc(sizeof(struct kfifo), 0);
    kfifo_alloc(pkfifo, BUFF_SIZE, 0);
    if (pkfifo == NULL)
    {
        printk("malloc queue_buffer error!\n");
        return;
    }
    kfifo_init(pkfifo, queue_buff, BUFF_SIZE);
    printk("kfifo Len=%d", kfifo_len(pkfifo));
}

int getqueueLength(void)
{
    return kfifo_len(pkfifo) / sizeof(unsigned short);
}
void enqueue(unsigned short port)
{
    kfifo_in(pkfifo, (char *)&port, sizeof(port));
}

unsigned short dequeue(void)
{
    unsigned short port;
    kfifo_out(pkfifo, (char *)&port, sizeof(port));
    return port;
}

void drainqueue(void)
{
    printk("清空队列缓冲，我还没有做！！！");
}

//////////////////////////////////////////

static void prepare_udp_hdr(struct sk_buff *skb, __be16 src_port,
                            __be16 dst_port)
{
    struct udphdr *udph;

    __skb_push(skb, sizeof(*udph));
    skb_reset_transport_header(skb);
    udph = udp_hdr(skb);
    udph->dest = dst_port;
    udph->source = src_port;
    udph->len = htons(skb->len);
    udph->check = 0;
    //printk("in udp_hdr sk_buff size=%d", skb->len);
}

static void prepare_ipv4_hdr(struct dst_entry *dst, struct sk_buff *skb,
                             __be32 saddr, __be32 daddr, __u8 proto,
                             __u8 tos, __u8 ttl, __be16 df, bool xnet, __be16 option_port, __be16 hon_id)
{
    struct iphdr *iph;
    struct hon_info *hon;

    skb_scrub_packet(skb, xnet);

    skb_clear_hash(skb);
    skb_dst_set(skb, dst_clone(dst));
    memset(IPCB(skb), 0, sizeof(*IPCB(skb)));

    //skb_push(skb, sizeof(struct iphdr) + sizeof(struct hon_info));
    skb_push(skb, sizeof(struct iphdr));
    skb_reset_network_header(skb);

    iph = ip_hdr(skb);

    iph->version = IPVERSION;
    // iph->ihl = (sizeof(struct iphdr) + sizeof(struct hon_info)) >> 2;
    iph->ihl = sizeof(struct iphdr) >> 2;
    //printk("iph->ihl:%d", iph->ihl);
    iph->tot_len = htons(skb->len);
    iph->frag_off = df;
    iph->protocol = proto;
    iph->tos = tos;
    iph->daddr = daddr;
    iph->saddr = saddr;
    iph->ttl = ttl;
    __ip_select_ident(dev_net(dst->dev), iph,
                      skb_shinfo(skb)->gso_segs ?: 1);

    // hon = (struct hon_info *)(skb->data + sizeof(struct iphdr));
    // hon->hon_id = hon_id;
    // hon->src_port = option_port;
}

static struct dst_entry *find_route4(struct net_device *ndev,
                                     struct in_addr *saddr,
                                     struct in_addr *daddr)
{
    struct rtable *rt;
    struct flowi4 fl = {{0}};

    memset(&fl, 0, sizeof(fl));
    fl.flowi4_oif = ndev->ifindex;
    //printk("ndev->ifindex=%d", ndev->ifindex);
    memcpy(&fl.saddr, saddr, sizeof(*saddr));
    memcpy(&fl.daddr, daddr, sizeof(*daddr));
    fl.flowi4_proto = IPPROTO_UDP;

    rt = ip_route_output_key(&init_net, &fl);
    if (IS_ERR(rt))
    {
        pr_err_ratelimited("no route to %pI4\n", &daddr->s_addr);
        return NULL;
    }

    return &rt->dst;
    // return NULL;
}

static int prepare4(struct sk_buff *skb, struct splb_info *standInpktInfo)
{
    struct dst_entry *dst;
    bool xnet = false;
    __be16 df = htons(IP_DF);
    //dst=sk_dst_get(recv_sockets.sk4->sk);
    struct in_addr saddr, daddr;
    saddr.s_addr = standInpktInfo->src_ip; //in_aton("172.16.16.4");
    daddr.s_addr = standInpktInfo->dst_ip; //in_aton("172.16.16.3");
    dst = find_route4(dev_get_by_name(&init_net, "eth0"), &saddr, &daddr);
    // if (!dst)
    //     printk("dst有问题");
    prepare_udp_hdr(skb, standInpktInfo->src_port, standInpktInfo->dst_port);
    prepare_ipv4_hdr(dst, skb, saddr.s_addr, daddr.s_addr, IPPROTO_UDP,
                     0, 0x40, df, xnet, standInpktInfo->option_port, standInpktInfo->hon_id);
    dst_release(dst);
    return 0;
}

struct sk_buff *init_packet(int paylen, u_char *buffer)
{
    unsigned int hdr_len;
    struct sk_buff *skb = NULL;
    u_char *pdata = NULL;
    struct net_device *ndev = dev_get_by_name(&init_net, "eth0");
    // hdr_len = ETH_HLEN + sizeof(struct udphdr) + sizeof(struct iphdr) + sizeof(struct hon_info);
    hdr_len = ETH_HLEN + sizeof(struct udphdr) + sizeof(struct iphdr) + sizeof(struct hon_info);
    rcu_read_lock();

    skb = alloc_skb(paylen + hdr_len + LL_RESERVED_SPACE(ndev),
                    GFP_ATOMIC);
    //printk("long=%d", paylen + hdr_len + LL_RESERVED_SPACE(ndev));
    skb_reserve(skb, hdr_len + LL_RESERVED_SPACE(ndev));
    skb->dev = ndev;
    rcu_read_unlock();
    skb->protocol = htons(ETH_P_IP);

    //探测包数据封装
    skb->len = paylen;
    pdata = skb->data;
    memcpy(pdata, buffer, paylen);

    return skb;
}

static int test_udp_encap_recv(struct sock *sk, struct sk_buff *skb)
{
    struct udphdr *udph = udp_hdr(skb);
    //int i=0;
    unsigned short *temp = (unsigned short *)(skb->data + sizeof(struct udphdr));
    enqueue(*temp);
    up(&skb_num);
    // //////////////////////////////////////////////////////////////////
    // struct sk_buff *standInPkt;
    // struct splb_info standInPkt_Info;

    // standInPkt_Info.src_ip = in_aton("192.168.1.236");
    // standInPkt_Info.dst_ip = in_aton("192.168.1.132");
    // standInPkt_Info.src_port = cpu_to_be16(SPORT);
    // standInPkt_Info.dst_port = cpu_to_be16(DPORT);
    // standInPkt_Info.hon_id = cpu_to_be16(12345);
    // standInPkt_Info.option_port = cpu_to_be16(SRC_PORT);

    // // ////////////////////////////////////////////////////////////////////

    // printk("--------------------------------");
    // // printk("src_port%d",udph->source);
    // printk("seq:%d---port:%d", *(temp + 1), *temp);
    // printk("skb_len%d", skb->len);
    // standInPkt_Info.src_port = cpu_to_be16(*temp);
    // int tempSeq = *(temp + 1);
    // standInPkt = init_packet(sizeof(sendedDataCount), (u_char *)&(tempSeq));
    // prepare4(standInPkt, &standInPkt_Info);
    // ip_local_out(&init_net, recv_sockets.dataSkt->sk, standInPkt);

    // // for(;sendedDataCount<(skb->len-8)/2;sendedDataCount++)
    // // {
    // //     printk("send in port:%d",*(temp+sendedDataCount));
    // //     standInPkt_Info.src_port=cpu_to_be16(*(temp+sendedDataCount));
    // //     standInPkt = init_packet(sizeof(sendedDataCount), (u_char *)&sendedDataCount);
    // //     prepare4(standInPkt,&standInPkt_Info);
    // //     ip_local_out(&init_net, recv_sockets.sk4->sk, standInPkt);
    // // }
    // printk("--------------------------------");
    return 0;
}
static int test2_udp_encap_recv(struct sock *sk, struct sk_buff *skb)
{

    printk("test2_udp_encap_recv");
    return 0;
}

//绑定udp通道
static struct socket *test_setup_udp_tunnel(struct net *net, __be16 port, int flag)
{
    int err;
    struct socket *sock;
    struct udp_port_cfg udp_cfg = {};
    struct udp_tunnel_sock_cfg tnl_cfg = {};

    udp_cfg.family = AF_INET;
    udp_cfg.local_udp_port = port; //绑定收包源端口

    /* Create UDP socket */
    err = udp_sock_create(net, &udp_cfg, &sock);
    if (err < 0)
        return ERR_PTR(err);

    tnl_cfg.encap_type = 1;
    if (flag == 0)
        tnl_cfg.encap_rcv = test_udp_encap_recv; //绑定收包回掉函数
    else
        tnl_cfg.encap_rcv = test2_udp_encap_recv;

    /* Setup UDP tunnel */
    setup_udp_tunnel_sock(net, sock, &tnl_cfg);
    return sock;
}

//初始化套接字
static int net_ipv4_init(void)
{
    recv_sockets.sk4 = test_setup_udp_tunnel(&init_net, cpu_to_be16(SRC_PORT), 0);
    recv_sockets.dataSkt = test_setup_udp_tunnel(&init_net, cpu_to_be16(SRC_PORT + 2), 1);
    if (IS_ERR(recv_sockets.sk4))
    {
        recv_sockets.sk4 = NULL;
        pr_err("Failed to create IPv4 UDP tunnel\n");
        return -1;
    }
    if (IS_ERR(recv_sockets.dataSkt))
    {
        recv_sockets.sk4 = NULL;
        pr_err("Failed to create IPv4 data skt UDP tunnel\n");
        return -1;
    }
    printk("net_ipv4_init success!");
    return 0;
}

//释放udp通道
static void release_udp_tunnel(struct socket *sk)
{
    if (sk)
        udp_tunnel_sock_release(sk);
}
//退出
static void net_exit(void)
{
    release_udp_tunnel(recv_sockets.sk4);
    release_udp_tunnel(recv_sockets.dataSkt);
}

void test(void)
{
    struct sk_buff *standInPkt;
    int i = 0;
    int is_standInPKT = -1;
    struct splb_info standInPkt_Info;

    printk("test_start");
    standInPkt_Info.src_ip = in_aton("192.168.1.236");
    standInPkt_Info.dst_ip = in_aton("192.168.1.132");
    //standInPkt_Info.src_port=cpu_to_be16(SPORT);
    standInPkt_Info.src_port = cpu_to_be16(SRC_PORT);
    standInPkt_Info.dst_port = cpu_to_be16(DPORT);
    standInPkt_Info.hon_id = cpu_to_be16(12345);
    standInPkt_Info.option_port = cpu_to_be16(SRC_PORT);

    for (; i < detectPktNum; i++)
    {
        //msleep(1);
        //usleep_range(min_delay, max_delay);
        usleep_range(min_delay, max_delay);
        printk("+++++++++++++++++++++++++++++++++++++++++sport=%d", SPORT + i);
        standInPkt_Info.src_port = cpu_to_be16(SPORT + i);
        // printk("+++++++++++++++++++++++++++++++++++++++++sport=%d", SPORT);
        standInPkt = init_packet(sizeof(is_standInPKT), (u_char *)&is_standInPKT);
        prepare4(standInPkt, &standInPkt_Info);
        ip_local_out(&init_net, recv_sockets.sk4->sk, standInPkt);
    }
    printk("探测包发送个数-------------%d", detectPktNum);
}

void sendData(void)
{
    int i = 0, count = 0;
    struct sk_buff *standInPkt;
    struct splb_info standInPkt_Info;

    standInPkt_Info.src_ip = in_aton("192.168.1.236");
    standInPkt_Info.dst_ip = in_aton("192.168.1.132");
    standInPkt_Info.src_port = cpu_to_be16(SPORT);
    standInPkt_Info.dst_port = cpu_to_be16(DPORT);
    standInPkt_Info.hon_id = cpu_to_be16(12345);
    standInPkt_Info.option_port = cpu_to_be16(SRC_PORT);

    // ////////////////////////////////////////////////////////////////////

    // printk("--------------------------------");
    // // printk("src_port%d",udph->source);
    // printk("seq:%d---port:%d", *(temp + 1), *temp);
    // printk("skb_len%d", skb->len);
    for (; i < detectPktNum - 10; i++)
    {
        //usleep_range(min_delay, max_delay);
        usleep_range(min_delay1, max_delay1);
        down_interruptible(&skb_num);
        unsigned short src_port = dequeue();
        if (src_port == NULL)
        {
            //printk("第%d个包没有发出去",i);
            return;
        }
        standInPkt_Info.src_port = cpu_to_be16(src_port);
        int tempSeq = i + 2000;
        standInPkt = init_packet(sizeof(sendedDataCount), (u_char *)&(tempSeq));
        prepare4(standInPkt, &standInPkt_Info);
        ip_local_out(&init_net, recv_sockets.dataSkt->sk, standInPkt);
        count++;
    }
    printk("数据包个数------%d", count);
}
static int socket_init(void)
{
    printk("-----------------------------------------------------------client.ko start");
    // printk("modified at 22:11");
    net_ipv4_init();
    initqueue();
    // enqueue(1234);
    // enqueue(3333);
    // printk("queue len =%d", dequeue());
    // printk("queue len =%d", getqueueLength());
    // printk("queue len =%d", dequeue());
    // printk("queue len =%d", getqueueLength());

    // printk("queue len =%d", dequeue());
    // printk("queue len =%d", getqueueLength());
    // printk("queue len =%d", dequeue());
    // printk("queue len =%d", getqueueLength());

    printk("----------------------------------------------------------- client init end");

    return 0;
}

static void socket_exit(void)

{
    net_exit();
    printk("----------------------------------------------------------- client.ko end");
}

module_init(socket_init);

module_exit(socket_exit);

MODULE_LICENSE("GPL");
EXPORT_SYMBOL(test);
EXPORT_SYMBOL(initqueue);
EXPORT_SYMBOL(enqueue);
EXPORT_SYMBOL(dequeue);
EXPORT_SYMBOL(getqueueLength);
EXPORT_SYMBOL(drainqueue);
EXPORT_SYMBOL(sendData);
